"""
文件工具方法
"""
import os

# import pypandoc
from docx2pdf import convert as convert_docx2pdf
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_unstructured import UnstructuredLoader
from langchain_community.document_loaders import (
    TextLoader,
    UnstructuredPDFLoader,
    PyPDFium2Loader,
    PyPDFLoader,
    CSVLoader,
    UnstructuredWordDocumentLoader,
    DirectoryLoader,
)
from exception.T2Exception import T2Exception



DOC_FORMAT = ['.pdf','.docx','.doc','.txt','.csv']


#校验文档格式
def check_file(file_path):
    name, ext = os.path.splitext(file_path)
    return ext.lower() in DOC_FORMAT

#文档分段
def doc_loader(file_path):
    # 加载、切分文档
    if check_file(file_path):
        if file_path.endswith(".pdf"):
            # UnstructuredPDFLoader，基于OCR，速度慢效果差
            # loader = UnstructuredPDFLoader(file_path=file_path,languages=["chi_sim"])
            # PyPDFium2Loader
            # loader = PyPDFium2Loader(file_path=file_path)
            # PyPDFLoader
            loader = PyPDFLoader(file_path=file_path)
        elif file_path.endswith(".docx") or file_path.endswith(".doc"):
            loader = UnstructuredWordDocumentLoader(file_path=file_path)
        elif file_path.endswith(".txt"):
            loader = TextLoader(file_path=file_path,encoding="utf-8")
        elif file_path.endswith(".csv"):
            loader = CSVLoader(file_path=file_path,encoding="utf-8")  ## 必须指定编码格式，不然load失败
        else:
            loader = UnstructuredLoader(file_path=file_path)
    else:
        raise T2Exception(f"文件格式不支持，{str(DOC_FORMAT)}");
    return loader

#文档分段
def doc_spliter(loader):
    # 加载、切分文档
    chunks = []
    texts = []
    if loader:
        data = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=750, chunk_overlap=50,separators=["\n\n"])
        chunks = text_splitter.split_documents(data)
    for chunk in chunks:
        texts.append(chunk.page_content)
    return texts

#docx 转 pdf
def docx_to_pdf(input_path, output_path):
    # pypandoc.convert_file(source_file=input_path, format='docx', outputfile=output_path, to="pdf")
    convert_docx2pdf(input_path, output_path)
     # 这个不支持linux
     # apt install libreoffice
     # libreoffice --convert-to pdf:writer_pdf_Export demo.docx --outdir ./

if __name__ == '__main__':
    docx_to_pdf("H:\\llmapp\\test_file\\动火作业报告.docx", "H:\\llmapp\\test_file\\动火作业报告.pdf")